from flask import Flask, request 
import json 
import cv2
import argparse

from src.engine_opencv import Engine 
from src.face_align import norm_crop 

from src.utils import get_logger
from src.grpc_client import GrpcClient
from src.global_dict import gd
import numpy as np
import urllib
import os
from ultralytics import YOLO

app = Flask(__name__)

@app.route('/face_api/add_source', methods=['POST'])
def add_source():
    logger.info('add source ====>')
    result = {'Code': 200,
              'Msg': 'add success!',
              'Data': {
                  'TaskList': []
              }}
    
    try:
        logger.info(f'input: {request.form}')

        task_id = request.form['task_id']
        stream_url = request.form['rtsp_url']
        # if stream_url == '0':
        #     stream_url = int(stream_url)
        grpc_address = request.form.get('grpc_address', None)
        decode_mode = request.form.get('decode_mode', 'cpu')

        skip_frame_rate = request.form.get('skip_frame_rate', 0)
        skip_frame_rate = int(skip_frame_rate)
        
        add_stranger = request.form.get('add_stranger', 1)
        add_stranger = int(add_stranger)
        alert_time = request.form.get('alert_time', 1)
        alert_time = int(alert_time)


        dbase_name = request.form.get('dbaseName', 'stream_test')

        rule_info = request.form.get('rule_info', {})
        
        if rule_info: rule_info = json.loads(rule_info)

     
        if not isinstance(rule_info, dict): rule_info = {}

        if task_id in engine.task_dict.keys():
            result['Code'] = 11001
            result['Msg'] = f'duplicate task id {task_id}'

            
        else:
            
            logger.info(f'add task task_id {task_id} ------')
            engine.add_source(task_id, stream_url, grpc_address, decode_mode, skip_frame_rate, add_stranger, alert_time, dbase_name,engine_file_path,rule_info)

        result['Data']['TaskList'] = list(engine.task_dict.keys())
        return result
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result
    

@app.route('/face_api/remove_source', methods=['DELETE'])
def remove_source():
    logger.info('remove source ====>')
    result = {'Code': 200,
              'Msg': 'remove success!',
              'Data':{
                  'TaskList': []
              }}
    
    try:
        logger.info(f'input: {request.form}')

        task_id = request.form['task_id']

        if task_id not in engine.task_dict.keys():
            result['Code'] = 11000
            result['Msg'] = f'can not find task id {task_id}'
        else:
            engine.remove_source(task_id)

        result['Data']['TaskList'] = list(engine.task_dict.keys())
        return result
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result

    
@app.route('/face_api/get_source_list', methods=['GET'])
def get_source_list():
    logger.info('get source list ====>')
    result = {'Code': 200,
              'Msg': 'get success!',
              'Data': {
                  'TaskList': list(engine.task_dict.keys())
              }}
    return result


@app.route('/face_api/source_status', methods=['POST'])
def source_status():
    logger.info('source_status ====>')
    result = {'Code': 200,
              'Msg': 'check success!',
              'Data': {}}
    
    try:
        logger.info(f'input: {request.get_json()}')
        task_ids = request.get_json()
        
        for task_id in task_ids:
            heart_beat = gd.heart_beat_dict.get(task_id, None) 
            if heart_beat is not None:
                result['Data'][task_id] = heart_beat.is_alive()
            else:
                result['Data'][task_id] = False

        logger.info(f'source status result: {result}')
        return result
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result


@app.route('/face_api/face_detect', methods=['POST'])
def face_detect():
    logger.info('face detect ====>')
    result = {'Code': 200,
              'Msg': 'response success!',
              'SceneCode': '',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'ImageWidth': 0,
                  'ImageHeight': 0,
                  'FaceNum': 0,
                  'FaceInfos': []
              }}

    for i in range(20):
        try:
            logger.info(f'input: {request.form}')

            img_url = request.form['image_url']
            dbase_name = request.form['dbaseName']
            face_recog = request.form.get('identify', 1)
            face_recog = int(face_recog)

            id_code = request.form['idCode']
            scene_code = request.form['sceneCode']

            result['IdCode'] = id_code
            result['SceneCode'] = scene_code

            cap = cv2.VideoCapture(img_url)
            _, img = cap.read(cv2.IMREAD_IGNORE_ORIENTATION)
            # res = urllib.request.urlopen(img_url)
            # img = np.asarray(bytearray(res.read()), dtype="uint8")
            # img = cv2.imdecode(img, cv2.IMREAD_COLOR)
            if img is None:
                result['Code'] = 5001
                result['Msg'] = f'img {img_url} is none'
            else:
                h, w, c = img.shape
                result['Data']['ImageWidth'] = w 
                result['Data']['ImageHeight'] = h 

                bboxes, kpss = engine.yolo_face.detect(img, args.face_detec_thres1)
                logger.info(f'face num in {img_url} is {len(bboxes)}')
                result['Data']['FaceNum'] = len(bboxes)

                for bbox, kps in zip(bboxes, kpss):
                    face_info = {}
                    face_info['Bbox'] = bbox.tolist()
                    face_info['Kps'] = kps.tolist()

                    if face_recog:
                        face = norm_crop(img, kps)

                        features, norms = engine.face_recog.get_features([face])
                        name, conf = engine.face_recog.top1(dbase_name, features[0])

                        print(name, conf)

                        if conf > 0:
                            face_info['FaceAttributesInfo'] = {}
                            face_info['FaceAttributesInfo']['FaceID'] = name
                            face_info['FaceAttributesInfo']['Score'] = conf
                    result['Data']['FaceInfos'].append(face_info)

            return result

        except Exception as e:
            logger.error(f'try error, {i}, {str(e)}')
            error_info = str(e)

    result['Code'] = 5002
    result['Msg'] = error_info
    return result

@app.route('/face_api/face_class', methods=['POST'])
def face_class():
    logger.info('face detect ====>')
    result = {'Code': 200,
              'Msg': 'response success!',
              'SceneCode': '',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'ImageWidth': 0,
                  'ImageHeight': 0,
                  'FaceNum': 0,
                  'FaceInfos': []
              }}

    for i in range(20):
        try:
            logger.info(f'input: {request.form}')

            img_url = request.form['image_url']
            dbase_name = request.form['dbaseName']
            face_recog = request.form.get('identify', 1)
            face_recog = int(face_recog)

            id_code = request.form['idCode']
            scene_code = request.form['sceneCode']

            result['IdCode'] = id_code
            result['SceneCode'] = scene_code


            for filename in os.listdir(img_url):
                face = cv2.imread(img_url+filename)

                

                features, norms = engine.face_recog.get_features([face])
                name, conf = engine.face_recog.top1(dbase_name, features[0])

                print(name, conf)

                   

            return result

        except Exception as e:
            logger.error(f'try error, {i}, {str(e)}')
            error_info = str(e)

    result['Code'] = 5002
    result['Msg'] = error_info
    return result

@app.route('/face_api/add_face_batch', methods=['POST'])
def add_face_batch():
    logger.info('add face ====>')
    overall_result = {'Code': 200,
                      'Msg': 'Batch add success!',
                      'Details': []}
    
    dbase_name = request.form['dbaseName']
    images_folder = request.form['images_folder']

    for filename in os.listdir(images_folder):
        if filename.endswith(('.jpg', '.png')):  # 根据需要调整格式
            result = {'RequestId': '',
                      'IdCode': '',
                      'Data': {
                          'FaceID': '',
                          'DbaseName': ''
                      }}
            image_path = os.path.join(images_folder, filename)
            face_id = os.path.splitext(filename)[0]  # 使用图片名作为faceID
            id_code = f'request_{filename}'  # 可以根据实际情况修改
            
            result['Data']['FaceID'] = face_id 
            result['Data']['DbaseName'] = dbase_name
            result['IdCode'] = id_code
            
            cap = cv2.VideoCapture(image_path)
            _, img = cap.read(cv2.IMREAD_IGNORE_ORIENTATION)
            if img is None:
                result['Code'] = 5001
                result['Msg'] = f'img {image_path} could not be read'
            elif engine.face_database.has_face(dbase_name, face_id):
                result['Code'] = 5004
                result['Msg'] = f'face id {face_id} already in face database'
            else:
                bboxes, kpss = engine.yolo_face.detect(img, args.face_detec_thres1)
                logger.info(f'face num detected is {len(bboxes)} for {filename} ------')

                if len(bboxes) ==  1:
                    face = norm_crop(img, kpss[0])
                    feature = engine.face_recog.get_features([face])[0]

                    engine.face_database.add_face(dbase_name, face_id, feature)

                    logger.info(f'face {face_id} added to redis ------')
                else:
                    result['Code'] = 5003
                    result['Msg'] = f'face num detected is {len(bboxes)} for {filename}'
            
            overall_result['Details'].append(result)
    
    return overall_result

@app.route('/face_api/add_face', methods=['POST'])
def add_face():
    logger.info('add face ====>')
    result = {'Code': 200,
              'Msg': 'add success!',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'FaceID': '',
                  'DbaseName': ''
              }}

    for i in range(10):
        try:
            logger.info(f'input: {request.form}')
            img_url = request.form['image_url']
            dbase_name = request.form['dbaseName']
            face_id = request.form['faceID']
            id_code = request.form['idCode']

            result['Data']['FaceID'] = face_id 
            result['Data']['DbaseName'] = dbase_name
            result['IdCode'] = id_code

            cap = cv2.VideoCapture(img_url)
            _, img = cap.read(cv2.IMREAD_IGNORE_ORIENTATION)
            # res = urllib.request.urlopen(img_url)
            # img = np.asarray(bytearray(res.read()), dtype="uint8")
            # img = cv2.imdecode(img, cv2.IMREAD_COLOR)
            if img is None:
                result['Code'] = 5001
                result['Msg'] = f'img {img_url} is none'
            elif engine.face_database.has_face(dbase_name, face_id):
                result['Code'] = 5004
                result['Msg'] = f'face id {face_id} already in face database'
            else:
                bboxes, kpss = engine.yolo_face.detect(img, args.face_detec_thres1)
                logger.info(f'face num detected is {len(bboxes)} ------')

                if len(bboxes) ==  1:
                    face = norm_crop(img, kpss[0])
                    feature = engine.face_recog.get_features([face])[0]

                    engine.face_database.add_face(dbase_name, face_id, feature)

                    logger.info(f'face {face_id} added to redis ------')
                else:
                    result['Code'] = 5003
                    result['Msg'] = f'face num detected is {len(bboxes)}'
       
            return result
        except Exception as e:
            logger.error(f'try error, {i}, {str(e)}') 
            error_info = str(e)

    result['Code'] = 5002
    result['Msg'] = error_info
    return result


@app.route('/face_api/remove_face', methods=['DELETE'])
def remove_face():
    logger.info('remove face ====>')
    result = {'Code': 200,
              'Msg': 'remove success!',
              'RequestId': '',
              'Data': {
                  'FaceID': '',
                  'DbaseName': ''
              }}
    try:
        logger.info(f'input: {request.form}')

        dbase_name = request.form['dbaseName']
        face_id = request.form['faceID']
        id_code = request.form['idCode']

        result['Data']['FaceID'] = face_id
        result['Data']['DbaseName'] = dbase_name
        result['IdCode'] = id_code 

        if not engine.face_database.has_face(dbase_name, face_id):
            result['Code'] = 501
            result['Msg'] = f'face_name {face_id} not in face list'
        else:
            engine.face_database.del_face(dbase_name, face_id)

        return result
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result


@app.route('/face_api/check_face', methods=['POST'])
def check_face():
    logger.info('check face ====>')
    result = {'Code': 200,
              'Msg': 'check success!',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'FaceID': '',
                  'DbaseName': '',
                  'Result': ''
              }}
    
    try:
        logger.info(f'input: {request.form}')

        dbase_name = request.form['dbaseName']
        face_id = request.form['faceID']
        id_code = request.form['idCode']

        result['Data']['FaceID'] = face_id
        result['Data']['DbaseName'] = dbase_name
        result['IdCode'] = id_code 

        if engine.face_database.has_face(dbase_name, face_id):
            result['Data']['Result'] = 'Exist'
        else:
            result['Data']['Result'] = 'Not exist'

        return result
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result


@app.route('/face_api/add_database', methods=['POST'])
def add_database():
    logger.info('add database ====>')
    result = {'Code': 200,
              'Msg': 'add success!',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'DbaseName': ''
              }}
    
    try:
        logger.info(f'input: {request.form}')

        dbase_name = request.form['dbaseName']
        id_code = request.form['idCode']

        result['Data']['DbaseName'] = dbase_name 
        result['IdCode'] = id_code

        if engine.face_database.has_db(dbase_name):
            result['Code'] = 5006
            result['Msg'] = f'database {dbase_name} already exists!'
        else:
            engine.face_database.add_db(dbase_name)
        
        return result
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result


@app.route('/face_api/remove_database', methods=['DELETE'])
def remove_database():
    logger.info('remove database ====>')
    result = {'Code': 200,
              'Msg': 'remove success!',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'DbaseName': ''
              }}
    
    try:
        logger.info(f'input: {request.form}')

        dbase_name = request.form['dbaseName']
        id_code = request.form['idCode']

        result['Data']['DbaseName'] = dbase_name 
        result['IdCode'] = id_code 
        
        if not engine.face_database.has_db(dbase_name):
            result['Code'] = 5008 
            result['Msg'] = f'database {dbase_name} does not exists!'
        else:
            engine.face_database.del_db(dbase_name)

        return result 
    except Exception as e:
        logger.error(f'try error, {str(e)}')
        result['Code'] = 5002
        result['Msg'] = str(e)
        return result 

###########################################################
@app.route('/face_api/img_detect', methods=['POST'])
def img_detect():
    logger.info('face detect ====>')
    result = {'Code': 200,
              'Msg': 'response success!',
              'SceneCode': '',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'ImageWidth': 0,
                  'ImageHeight': 0,
                  'PersonNum': 0,
                  'PersonInfos': []
              }}

    for i in range(20):
        try:
            logger.info(f'input: {request.form}')

            img_url = request.form['image_url']
            dbase_name = request.form['dbaseName']
            face_recog = request.form.get('identify', 1)
            face_recog = int(face_recog)

            id_code = request.form['idCode']
            scene_code = request.form['sceneCode']

            result['IdCode'] = id_code
            result['SceneCode'] = scene_code
            

            # cap = cv2.VideoCapture(img_url)
            # _, img = cap.read(cv2.IMREAD_IGNORE_ORIENTATION)
            # res = urllib.request.urlopen(img_url)
            # img = np.asarray(bytearray(res.read()), dtype="uint8")
            # img = cv2.imdecode(img, cv2.IMREAD_COLOR)

            img = cv2.imread(img_url)
            if img is None:
                result['Code'] = 5001
                result['Msg'] = f'img {img_url} is none'
            else:
                h, w, c = img.shape
                result['Data']['ImageWidth'] = w 
                result['Data']['ImageHeight'] = h 
                
                person_results = yolo11n_pose_img(img,imgsz=960, device=0, iou=0.7, conf=0.25, verbose=False)  # predict on an image
                logger.info(f'person num in {img_url} is {len(person_results[0])}')
                result['Data']['PersonNum'] = len(person_results[0])
                for res in person_results[0]:
                    if res.boxes is None:
                        continue

                    person_info = {}

                    xyxy = res.boxes.xyxy.cpu().numpy().tolist()[0]
                    xy = res.keypoints.xy.cpu().numpy().tolist()[0]
                    
                    x1, y1, x2, y2 = map(int, xyxy)

                    foot_point = [[round(x, 2) for x in xy[16]], [round(x, 2) for x in xy[15]]]
                    # foot_point = [xy[16],xy[15]]
                    nose_point = [round(x, 2) for x in xy[0]]

                    img_bgr_body = img[y1:y2, x1:x2].copy()
                    
                    person_info['footPoint'] = foot_point
                    person_info['personBox'] = [x1, y1, x2, y2]

                    bboxes, kpss = engine.yolo_face.detect(img_bgr_body, args.face_detec_thres1)
                    
                    name='stranger'
                    for bbox, kps in zip(bboxes, kpss):
                        
                        if is_point_in_bbox(nose_point, bbox[0:4],[x1, y1]):
                            face = norm_crop(img_bgr_body, kps)

                            features, norms = engine.face_recog.get_features([face])
                            name, conf = engine.face_recog.top1(dbase_name, features[0],'stranger')

                           
                    print(name)
                    
                    
                    person_info['faceID'] = name
            
                    result['Data']['PersonInfos'].append(person_info)

                data = result['Data']['PersonInfos']

                drow_mag(img, data)

                return result

        except Exception as e:
            logger.error(f'try error, {i}, {str(e)}')
            error_info = str(e)

    result['Code'] = 5002
    result['Msg'] = error_info
    return result

def drow_mag(img,data):


    # 设置字体和颜色
    font = cv2.FONT_HERSHEY_SIMPLEX
    box_color = (0, 255, 0)  # 绿色框
    point_color = (0, 0, 255)  # 红色点
    text_color = (255, 255, 255)  # 白色文字

    for item in data:
        face_id = item['faceID']
        box = item['personBox']  # [x1, y1, x2, y2]
        foot_points = item['footPoint']  # [[x1,y1], [x2,y2]]

        # 绘制 personBox
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(img, (x1, y1), (x2, y2), box_color, 2)

        # 显示 faceID
        cv2.putText(img, face_id, (x1, y1 + 10), font, 1, text_color, 3)

        # 绘制 footPoint 上的两个点
        for fx, fy in foot_points:
            if fx > 0 and fy > 0:  # 忽略 [0,0] 这类无效点
                cv2.circle(img, (int(fx), int(fy)), radius=5, color=point_color, thickness=-1)

    # 保存结果
    cv2.imwrite('im.jpg', img)
def is_point_in_bbox(point, box, ad):
    x, y = point
    x1, y1, x2, y2 = box
    return x1+ad[0] <= x <= x2+ad[0] and y1+ad[1] <= y <= y2+ad[1]


@app.route('/face_api/img_detect_violation', methods=['POST'])
def img_detect_violation():
    logger.info('face detect ====>')
    result = {'Code': 200,
              'Msg': 'response success!',
              'SceneCode': '',
              'RequestId': '',
              'IdCode': '',
              'Data': {
                  'ImageWidth': 0,
                  'ImageHeight': 0,
                  'PersonNum': 0,
                  'PersonInfos': []
              }}

    for i in range(20):
        try:
            logger.info(f'input: {request.form}')

            img_url = request.form['image_url']
            dbase_name = request.form['dbaseName']
            face_recog = request.form.get('identify', 1)
            face_recog = int(face_recog)

            id_code = request.form['idCode']
            scene_code = request.form['sceneCode']

            result['IdCode'] = id_code
            result['SceneCode'] = scene_code
            

            # cap = cv2.VideoCapture(img_url)
            # _, img = cap.read(cv2.IMREAD_IGNORE_ORIENTATION)
            # res = urllib.request.urlopen(img_url)
            # img = np.asarray(bytearray(res.read()), dtype="uint8")
            # img = cv2.imdecode(img, cv2.IMREAD_COLOR)

            img = cv2.imread(img_url)
            if img is None:
                result['Code'] = 5001
                result['Msg'] = f'img {img_url} is none'
            else:
                h, w, c = img.shape
                result['Data']['ImageWidth'] = w 
                result['Data']['ImageHeight'] = h 
                
                person_seg_results = yolo11n_seg_img(img,imgsz=960, device=0, iou=0.7, conf=0.25, verbose=False)  # predict on an image

                logger.info(f'person num in {img_url} is {len(person_results[0])}')
                result['Data']['PersonNum'] = len(person_results[0])
                for res in person_results[0]:
                    if res.boxes is None:
                        continue

                    person_info = {}

                    xyxy = res.boxes.xyxy.cpu().numpy().tolist()[0]
                    xy = res.keypoints.xy.cpu().numpy().tolist()[0]
                    
                    x1, y1, x2, y2 = map(int, xyxy)

                    foot_point = [[round(x, 2) for x in xy[16]], [round(x, 2) for x in xy[15]]]
                    # foot_point = [xy[16],xy[15]]
                    nose_point = [round(x, 2) for x in xy[0]]

                    img_bgr_body = img[y1:y2, x1:x2].copy()
                    
                    person_info['footPoint'] = foot_point
                    person_info['personBox'] = [x1, y1, x2, y2]

                    bboxes, kpss = engine.yolo_face.detect(img_bgr_body, args.face_detec_thres1)
                    
                    name='stranger'
                    for bbox, kps in zip(bboxes, kpss):
                        
                        if is_point_in_bbox(nose_point, bbox[0:4],[x1, y1]):
                            face = norm_crop(img_bgr_body, kps)

                            features, norms = engine.face_recog.get_features([face])
                            name, conf = engine.face_recog.top1(dbase_name, features[0],'stranger')

                           
                    print(name)
                    
                    
                    person_info['faceID'] = name
            
                    result['Data']['PersonInfos'].append(person_info)

                data = result['Data']['PersonInfos']

                drow_mag(img, data)

                return result

        except Exception as e:
            logger.error(f'try error, {i}, {str(e)}')
            error_info = str(e)

    result['Code'] = 5002
    result['Msg'] = error_info
    return result





def get_args():
    parser = argparse.ArgumentParser(description='server args')
    parser.add_argument('--triton_port', type=int, default=8111)
    parser.add_argument('--face_detec_model_name', type=str, default='ensemble_yolov8_face')
    parser.add_argument('--face_recog_model_name', type=str, default='face_recog')
    parser.add_argument('--face_detec_thres1', type=float, default=0.15)
    parser.add_argument('--face_detec_thres2', type=float, default=0.15)
    parser.add_argument('--redis_port', type=int, default=8110)
    parser.add_argument('--face_database_thres', type=float, default=0.25)
    parser.add_argument('--server_port', type=int, default=6000, help='Port to run the Flask server on')


    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    logger = get_logger('./logs')

    engine_file_path = "model2/yolo11s.onnx"
    yolo_pose_path = "model2/yolo11m-pose.pt"
    yolo_seg_path = "model2/yolo11n-seg.pt"

    yolo11n_pose_img = YOLO(yolo_pose_path)
    yolo11n_seg_img = YOLO(yolo_seg_path)
    

    engine = Engine(logger, args)

    logger.info('start run server ------')
    app.run(host='0.0.0.0', port=args.server_port)
